# encoding=utf-8
import time
import cv2
import numpy
import openvino as ov


class DensenetGrade:
    def __init__(self, pth_path: str):
        core = ov.Core()
        cpu_model = core.read_model(pth_path)
        self.model = core.compile_model(model=cpu_model, device_name="CPU")
        self.output_key = self.model.output(0)

    def softmax(self, x):
        f_x = numpy.exp(x) / numpy.sum(numpy.exp(x))
        return f_x

    def infer(self, image_path: str):
        trans_dim = self._preprocess(image_path)
        result = self.model([trans_dim])[self.output_key]
        y = self.softmax(result[0]).tolist()
        return self._postprocess(y)

    def _postprocess(self, pred_arr):
        return {
            'index': pred_arr.index(max(pred_arr)),
            'pie': pred_arr
        }

    def _preprocess(self, data_in: str):
        img_arr = self.circle_crop(data_in) / 255.0
        x = numpy.expand_dims(img_arr[:, :, ::-1].astype(numpy.float32).transpose((2, 0, 1)), axis=0)
        return x

    def circle_crop(self, image_src: str):
        crop_mask = self.crop_image_from_mask(image_src)
        return self.crop_image_with_gaussian(crop_mask)

    def crop_image_from_mask(self, image_src: str):
        # load
        image = cv2.imread(image_src)

        # binary
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, binary_image = cv2.threshold(gray_image, 7, 255, cv2.THRESH_BINARY)

        # cal roi
        x, y, w, h = cv2.boundingRect(binary_image)
        center = (w // 2), (h // 2)
        radius = min(center)
        y = y + center[1] - radius
        x = x + center[0] - radius
        copy_image = image[y: y + 2 * radius, x: x + 2 * radius]

        # gen mask
        mask = numpy.zeros_like(copy_image)
        cv2.circle(mask, (radius, radius), radius, (1, 1, 1), -1)

        # exposure
        return copy_image * mask

    def crop_image_with_gaussian(self, data_in: numpy.ndarray):
        ori_image = cv2.resize(data_in, (224, 224)).astype(numpy.float32)
        out = cv2.GaussianBlur(ori_image, (11, 11), 10.0)
        exposure = cv2.addWeighted(ori_image, 4, out, -4, 128)
        exposure = numpy.clip(exposure, 0, 255).astype(numpy.uint8)
        exposure = cv2.cvtColor(exposure, cv2.COLOR_BGR2RGB)
        return exposure


if __name__ == '__main__':
    import os
    CONFIG_PATH = os.path.join(os.path.dirname(__file__))
    MODEL_PATH = os.path.join(CONFIG_PATH, "export_dense121_cpu.xml")
    TEST_PATH = os.path.join(CONFIG_PATH, "1.jpg")
    OUT1 = os.path.join(CONFIG_PATH, "2-1.jpg")
    OUT2 = os.path.join(CONFIG_PATH, "2-2.jpg")
    print(MODEL_PATH, TEST_PATH)
    # {'index': 2, 'pie': [0.14742931723594666, 0.06031818687915802, 0.6965796947479248, 0.09067228436470032, 0.005000528413802385]}
    grade = DensenetGrade(MODEL_PATH)
    # test
    cv2.imwrite(OUT1, grade.crop_image_from_mask(TEST_PATH))
    cv2.imwrite(OUT2, grade.circle_crop(TEST_PATH))
    # time cost
    t1 = time.perf_counter()
    print(grade.infer(TEST_PATH))
    t2 = time.perf_counter()
    print(t2 - t1)
